use concurrent_queue::ConcurrentQueue;
use ndarray::{
    Array, ArrayBase, ArrayView, Data, DataMut, IntoDimension, IxDyn, RawData, ShapeBuilder,
};
use num_complex::Complex64;
use std::{
    borrow::Cow,
    sync::{Arc, LazyLock, Mutex},
};
pub use tonglun_thread::get_current_thread_id;

static OPAQUE_VTABLES: boxcar::Vec<OpaqueVTable> = boxcar::Vec::new();
static RUNTIME_DESTRUCTION_QUEUE: LazyLock<
    dashmap::DashMap<RuntimeIdent, Arc<ConcurrentQueue<ForeignDestructor>>>,
> = LazyLock::new(|| dashmap::DashMap::new());
static RUNTIMES: boxcar::Vec<Runtime> = boxcar::Vec::new();

static SHARED_FUNCTIONS: LazyLock<dashmap::DashMap<String, Arc<Function>>> =
    LazyLock::new(|| dashmap::DashMap::new());

static ERRORS: Mutex<Vec<ErrorMsgData>> = Mutex::new(Vec::new());
static ERROR_FREE_LIST: Mutex<Vec<usize>> = Mutex::new(Vec::new());

#[derive(Clone)]
#[repr(C)]
pub struct Function {
    name: String,
    runtime: RuntimeIdent,
    foreign_destructor: ForeignDestructor,
    fptr: unsafe extern "C" fn(
        out: *mut CDatum,
        closure: *mut u8,
        args: *const CDatum,
        nargs: usize,
    ) -> ErrorMsg,
}

#[derive(Clone)]
pub struct OpaqueData {
    vtable: u32,
    data: *mut u8,
}

impl Drop for OpaqueData {
    fn drop(&mut self) {
        if let Some(vtable) = OPAQUE_VTABLES.get(self.vtable as usize) {
            let destructor = vtable.foreign_destructor;
            let destruction = ForeignDestructor {
                drop: destructor,
                closure: self.data as *mut u8,
            };
            drop_foreign(destruction, vtable.runtime, None);
        }
    }
}

impl std::fmt::Debug for OpaqueData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let vtable = OPAQUE_VTABLES.get(self.vtable as usize).unwrap();
        let name = vtable.name.as_str();
        write!(f, "OpaqueData(class={}, handle={:?})", name, self.data)
    }
}

impl std::fmt::Debug for Function {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let runtime_name = self.runtime.try_get_name().unwrap_or("unknown");
        write!(
            f,
            "Function(rt={}, name={})",
            runtime_name,
            self.name.as_str()
        )
    }
}

/// OpaqueVTable stores metadata for the opaque data.
///
/// It is used for registering the runtime-specific objects (such as classes in Python)
/// to the core.
///
/// XXX: understanding common polymorphisms using this?
#[repr(C)]
#[derive(Clone)]
struct OpaqueVTable {
    name: String,
    foreign_destructor: unsafe extern "C" fn(closure: *mut u8),
    runtime: RuntimeIdent,
}

pub struct Runtime {
    thread_safe: bool,
    name: String,
    def_thread: u64,
}

// this is used for registering the runtime to the core;
// designed for compatibility with downstream frameworks
#[repr(C)]
struct CRuntime {
    name: *const u8,
    thread_safe: u8,
}

unsafe impl Send for OpaqueVTable {}
unsafe impl Sync for OpaqueVTable {}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct RuntimeIdent(usize);

impl RuntimeIdent {
    pub fn try_get_name(&self) -> Option<&'static str> {
        RUNTIMES.get(self.0).map(|s| s.name.as_str())
    }
    pub fn try_get_thread_safe(&self) -> Option<bool> {
        RUNTIMES.get(self.0).map(|s| s.thread_safe)
    }
    pub fn get_def_thread(&self) -> Option<u64> {
        RUNTIMES.get(self.0).map(|s| s.def_thread)
    }

    pub fn get_runtime(&self) -> Option<&'static Runtime> {
        RUNTIMES.get(self.0)
    }
}

impl Function {
    pub fn invoke(&self, out: *mut CDatum, args: *const CDatum, nargs: usize) -> ErrorMsg {
        if self.runtime.get_def_thread() != Some(get_current_thread_id()) {
            let runtime_name = self.runtime.try_get_name().unwrap_or("unknown");
            let function_name = self.name.as_str();
            return ErrorMsg::create_error(
                format!(
                    "function ({}) from runtime ({}) is not in the definition thread",
                    runtime_name, function_name
                )
                .into(),
            );
        }
        unsafe { (self.fptr)(out, self.foreign_destructor.closure, args, nargs) }
    }
}

impl Drop for Function {
    fn drop(&mut self) {
        drop_foreign(self.foreign_destructor, self.runtime, None);
    }
}

fn run_finalizers(runtime: RuntimeIdent) -> ErrorMsg {
    let queue = if let Some(queue) = RUNTIME_DESTRUCTION_QUEUE.get(&runtime).map(|q| q.clone()) {
        queue
    } else {
        let name = runtime.try_get_name().unwrap_or("unknown");
        return ErrorMsg::create_error(
            format!("runtime ({}) has no destruction queue", name).into(),
        );
    };

    let def_thread = runtime.get_def_thread().unwrap();
    if def_thread != get_current_thread_id() {
        let name = runtime.try_get_name().unwrap_or("unknown");
        return ErrorMsg::create_error(
            format!(
                "running finalizers for runtime ({}) in a thread other than the definition thread",
                name
            )
            .into(),
        );
    }

    if queue.len() > 0 {
        while let Ok(destructor) = queue.pop() {
            // the downstream developer should ensure that the destructor does not crash
            unsafe { (destructor.drop)(destructor.closure) };
        }
    }

    ErrorMsg::ok()
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum Tag {
    I64,
    F64,
    Bool,
    U8Str,
    CF64,

    /// Vec<Ref>, HList/Tuple constructs in the backend languages
    HList,

    F64Tensor,
    CF64Tensor,
    I64Tensor,
    I32Tensor,

    Opaque,
    Function,
}

static MESSAGE_INVALID_TYPECODE: &str = "invalid typecode\0";
static MESSAGE_EMPTY_SHAPE: &str = "empty shape\0";
static MESSAGE_LIST_ALLOC_OUT_OF_MEMORY: &str = "list alloc out of memory\0";
static MESSAGE_TENSOR_ALLOC_OUT_OF_MEMORY: &str = "tensor alloc out of memory\0";
static MESSAGE_NOT_TENSOR: &str = "tensor operation not a tensor\0";
static MESSAGE_OUT_OF_BOUNDS: &str = "out of bounds\0";
static MESSAGE_NOT_I64: &str = "not integer\0";
static MESSAGE_NOT_F64: &str = "not double float\0";
static MESSAGE_NOT_CF64: &str = "not complex float\0";
static MESSAGE_NOT_BOOL: &str = "not boolean\0";
static MESSAGE_NOT_U8STR: &str = "not string\0";
static MESSAGE_NOT_HLIST: &str = "not list\0";
static MESSAGE_UNKNOWN_OPAQUE_VTABLE: &str = "unknown opaque vtable\0";
static MESSAGE_NULL_REFERENCE: &str = "null reference\0";

#[repr(transparent)]
pub struct ErrorMsgData(Cow<'static, str>);

impl ErrorMsgData {
    pub fn new(message: Cow<'static, str>) -> Self {
        match message {
            Cow::Borrowed(message) => {
                if !message.ends_with('\0') {
                    let mut message = message.to_string();
                    message.push('\0');
                    return Self(Cow::Owned(message));
                }
                Self(Cow::Borrowed(message))
            }
            Cow::Owned(mut message) => {
                if !message.ends_with('\0') {
                    message.push('\0');
                }
                Self(Cow::Owned(message))
            }
        }
    }
}

#[repr(C)]
pub struct ErrorMsg(usize);

impl ErrorMsg {
    pub fn ok() -> Self {
        Self(0)
    }

    pub fn is_ok(&self) -> bool {
        self.0 == 0
    }

    pub fn create_error(message: Cow<'static, str>) -> Self {
        let error = ErrorMsgData::new(message);
        // try to get an id from the free list
        let id = {
            if let Some(id) = ERROR_FREE_LIST.lock().unwrap().pop() {
                Some(id)
            } else {
                None
            }
        };

        let id = if let Some(id) = id {
            ERRORS.lock().unwrap()[id - 1] = error;
            id
        } else {
            let mut errors = ERRORS.lock().unwrap();
            errors.push(error);
            errors.len()
        };

        Self(id)
    }

    pub unsafe fn get_error_msg_ptr(&self) -> *const u8 {
        if let Some(error_id) = self.0.checked_sub(1) {
            return ERRORS.lock().unwrap()[error_id]
                .0
                .as_ref()
                .as_bytes()
                .as_ptr();
        }
        return std::ptr::null();
    }

    pub fn clone_error_msg(&self) -> Option<String> {
        if self.is_ok() {
            return None;
        }
        let error_id = self.0 - 1;
        let lock = ERRORS.lock().unwrap();
        let error = lock.get(error_id)?;
        Some(error.0.as_ref().to_string())
    }

    pub unsafe fn free(&self) {
        if let Some(error_id) = self.0.checked_sub(1) {
            {
                let mut free_list = ERROR_FREE_LIST.lock().unwrap();
                free_list.push(error_id + 1);
            }
            {
                let mut errors = ERRORS.lock().unwrap();
                errors[error_id] = ErrorMsgData::new(Cow::Borrowed(""));
            }
        }
    }
}

/// CDatum does not own the data it points to.
/// It behaves like a reference, in order to ease the FFI manipulation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(C)]
pub struct CDatum {
    value: u64,
    extra: u32,
    tag: Tag,
}

#[derive(Debug)]
pub enum Tensor<T: 'static> {
    /// Native Tensor represented in Rust
    ///
    /// NOTE that the data is promised to be contiguous
    /// if it is created using valid constructors provided
    /// by this crate.
    Rust(Array<T, IxDyn>),
    /// Foreign Tensor represented in FFI
    /// It is also wrapped using Rust ndarray's ArrayView.
    ///
    /// NOTE that for the FFI case, the orginal tensor data
    /// is never modified once sent to the foreign language.
    FFI {
        array: ArrayView<'static, T, IxDyn>,
        destructor: ForeignDestructor,

        /// Overwrite the thread-safe flag of the runtime.
        ///
        /// This is useful when the runtime is not thread-safe
        /// but the tensor could be thread-safe.
        thread_safe: bool,
        /// Indicate the runtime of the tensor.
        ///
        /// This helps to distinguish the chance of running finalizers
        /// if the tensor is not thread-safe.
        runtime: RuntimeIdent,
    },
}

#[derive(Debug, Clone, Copy)]
#[repr(C)]
pub struct ForeignDestructor {
    drop: unsafe extern "C" fn(closure: *mut u8),
    closure: *mut u8,
}

unsafe impl Send for ForeignDestructor {}
unsafe impl Sync for ForeignDestructor {}

fn drop_foreign(destructor: ForeignDestructor, runtime: RuntimeIdent, thread_safe: Option<bool>) {
    let runtime_data = match runtime.get_runtime() {
        Some(runtime) => runtime,
        None => {
            println!("runtime not found");
            return;
        }
    };

    if thread_safe.unwrap_or_else(|| runtime_data.thread_safe)
        || get_current_thread_id() == runtime_data.def_thread
    {
        unsafe {
            (destructor.drop)(destructor.closure);
        }

        // TODO: automatically run finalizers for the runtime when some conditions are met
    } else {
        let queue = RUNTIME_DESTRUCTION_QUEUE
            .entry(runtime)
            .or_insert_with(|| Arc::new(concurrent_queue::ConcurrentQueue::unbounded()))
            .clone();

        match queue.push(destructor) {
            Ok(_) => {}
            Err(_) => {
                println!(
                    "failed to push destructor to queue (runtime = {})",
                    runtime_data.name
                );
            }
        }
    }
}

impl<T: Clone> Clone for Tensor<T> {
    fn clone(&self) -> Self {
        match self {
            Tensor::Rust(array) => Tensor::Rust(array.clone()),
            Tensor::FFI { array, .. } => {
                let array = array.to_owned();
                Tensor::Rust(array)
            }
        }
    }
}

pub fn linear_index_to_cartesian_index<'a, S>(
    array: &'a ArrayBase<S, IxDyn>,
    index: usize,
) -> Option<&'a S::Elem>
where
    S: RawData + Data,
{
    let order = array.as_slice_memory_order()?;
    let index = order.get(index)?;
    Some(index)
}

pub fn linear_index_to_cartesian_index_mut<'a, S>(
    array: &'a mut ArrayBase<S, IxDyn>,
    index: usize,
) -> Option<&'a mut S::Elem>
where
    S: RawData + DataMut,
{
    let order = array.as_slice_memory_order_mut()?;
    let index = order.get_mut(index)?;
    Some(index)
}

impl<T> Drop for Tensor<T> {
    fn drop(&mut self) {
        if let Tensor::FFI {
            destructor,
            thread_safe,
            runtime,
            ..
        } = self
        {
            drop_foreign(*destructor, *runtime, Some(*thread_safe));
        }
    }
}

impl<T: 'static + std::fmt::Debug> Tensor<T> {
    pub unsafe fn adapt_native(
        data: *const T,
        shape: &[usize],
        strides: Option<&[isize]>,
        drop: unsafe extern "C" fn(closure: *mut u8),
        closure: *mut u8,
        is_fortran_order: bool,
        thread_safe: bool,
        runtime: RuntimeIdent,
    ) -> Self {
        let empty_slice: &'static [T] = &[];
        let data = if data.is_null() {
            empty_slice.as_ptr()
        } else {
            data
        };

        if let Some(strides) = strides {
            // rust ndarray uses [usize] to store strides, but it is actually used as [isize].
            // ref: https://github.com/rust-ndarray/ndarray/blob/d5f32ec06e27d8705dcf5da4e6596abf51188909/src/dimension/axes.rs#L90
            let strides: &[usize] = unsafe { std::mem::transmute(strides) };

            let shape: Vec<usize> = shape.into();
            let shape = shape.set_f(is_fortran_order);

            // TODO: use numpy to test the validity of strides construction.
            let shape = shape.strides(strides.into_dimension());
            let array: ArrayView<'static, _, _> = unsafe { ArrayView::from_shape_ptr(shape, data) };
            return Tensor::FFI {
                array,
                destructor: ForeignDestructor { drop, closure },
                thread_safe,
                runtime,
            };
        } else {
            let shape: Vec<usize> = shape.into();
            let shape = shape.set_f(is_fortran_order);

            let array: ArrayView<'static, _, _> = unsafe { ArrayView::from_shape_ptr(shape, data) };
            return Tensor::FFI {
                array,
                destructor: ForeignDestructor { drop, closure },
                thread_safe,
                runtime,
            };
        }
    }

    pub fn numel(&self) -> usize {
        match self {
            Tensor::Rust(array) => array.len(),
            Tensor::FFI { array, .. } => array.len(),
        }
    }

    pub fn size(&self, axis: usize) -> usize {
        match self {
            Tensor::Rust(array) => {
                let dims = array.dim();
                let ndims = array.ndim();
                if axis >= ndims {
                    return 1;
                }
                dims[axis].into()
            }
            Tensor::FFI { array, .. } => {
                let dims = array.dim();
                let ndims = array.ndim();
                if axis >= ndims {
                    return 1;
                }
                dims[axis].into()
            }
        }
    }

    pub fn shape(&self) -> &[usize] {
        match self {
            Tensor::Rust(array) => array.shape(),
            Tensor::FFI { array, .. } => array.shape(),
        }
    }

    pub fn strides(&self) -> &[isize] {
        match self {
            Tensor::Rust(array) => array.strides(),
            Tensor::FFI { array, .. } => array.strides(),
        }
    }

    pub fn get_element(&self, index: usize) -> Option<&T> {
        match self {
            Tensor::Rust(array) => linear_index_to_cartesian_index(array, index),
            Tensor::FFI { array, .. } => linear_index_to_cartesian_index(array, index),
        }
    }

    pub fn get_data_ptr(&self) -> *const T {
        match self {
            Tensor::Rust(array) => array.as_ptr(),
            Tensor::FFI { array, .. } => array.as_ptr(),
        }
    }
}

impl<T: 'static + Default + Clone> Tensor<T> {
    pub fn prealloc(shape: &[usize], is_fortran_order: bool) -> Result<Self, ()> {
        let mut data = Vec::new();
        let len: usize = shape.iter().product();
        if let Err(_) = data.try_reserve(len) {
            return Err(());
        }

        // fill with default values
        data.resize(len, T::default());

        let shape = if is_fortran_order {
            shape.set_f(true)
        } else {
            shape.set_f(false)
        };

        // create the tensor with the shape
        match Array::from_shape_vec(shape, data) {
            Ok(array) => Ok(Tensor::Rust(array)),
            Err(_) => Err(()),
        }
    }
}

impl<T: 'static + Clone> Tensor<T> {
    pub fn get_element_mut(&mut self, index: usize) -> Option<&mut T> {
        match self {
            Tensor::Rust(array) => linear_index_to_cartesian_index_mut(array, index),
            Tensor::FFI { .. } => {
                *self = self.clone();
                if let Tensor::Rust(array) = self {
                    return linear_index_to_cartesian_index_mut(array, index);
                }
                panic!("fatal: FFI tensor should be converted to Rust tensor before mutating");
            }
        }
    }

    pub fn get_data_ptr_mut(&mut self) -> *mut T {
        match self {
            Tensor::Rust(array) => array.as_mut_ptr(),
            Tensor::FFI { .. } => {
                *self = self.clone();
                if let Tensor::Rust(array) = self {
                    return array.as_mut_ptr();
                }
                panic!("fatal: FFI tensor should be converted to Rust tensor before mutating");
            }
        }
    }
}

#[derive(Clone, Debug)]
pub enum Datum {
    I64(i64),
    F64(f64),
    Bool(u8),
    U8Str(Arc<String>),
    CF64(Arc<Complex64>),
    HList(Arc<Vec<Datum>>),

    F64Tensor(Arc<Tensor<f64>>),
    CF64Tensor(Arc<Tensor<Complex64>>),
    I64Tensor(Arc<Tensor<i64>>),
    I32Tensor(Arc<Tensor<i32>>),
    Opaque(Arc<OpaqueData>),
    Function(Arc<Function>),
}
struct MoveOrClonedArray {
    cloned: bool,
    variant: MoveOrClonedArrayVariant,
}

enum MoveOrClonedArrayVariant {
    F64Tensor(Arc<Tensor<f64>>),
    CF64Tensor(Arc<Tensor<Complex64>>),
    I64Tensor(Arc<Tensor<i64>>),
    I32Tensor(Arc<Tensor<i32>>),
}

impl Datum {
    pub unsafe fn drop_reference(value: CDatum) {
        match value.tag {
            Tag::CF64 => {
                let ptr: *const Complex64 = value.value as *const Complex64;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::U8Str => {
                let ptr: *const String = value.value as *const String;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::HList => {
                let ptr: *const Vec<Datum> = value.value as *const Vec<Datum>;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::F64Tensor => {
                let ptr: *const Tensor<f64> = value.value as *const Tensor<f64>;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::CF64Tensor => {
                let ptr: *const Tensor<Complex64> = value.value as *const Tensor<Complex64>;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::I64Tensor => {
                let ptr: *const Tensor<i64> = value.value as *const Tensor<i64>;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::I32Tensor => {
                let ptr: *const Tensor<i32> = value.value as *const Tensor<i32>;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::Opaque => {
                let ptr: *const OpaqueData = value.value as *const OpaqueData;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::Function => {
                let ptr = value.value as *const Function;
                let _ = unsafe { Arc::from_raw(ptr) };
            }
            Tag::I64 | Tag::F64 | Tag::Bool => { /* value type, do nothing */ }
        }
    }
    #[inline]
    pub fn create_owner(value: CDatum) -> Self {
        match value.tag {
            Tag::I64 => Datum::I64(unsafe { std::mem::transmute(value.value) }),
            Tag::F64 => Datum::F64(unsafe { std::mem::transmute(value.value) }),
            Tag::Bool => Datum::Bool(if value.value != 0 { 1 } else { 0 }),
            Tag::U8Str => {
                let ptr: *const String = value.value as *const String;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::U8Str(cloned)
            }
            Tag::CF64 => {
                let ptr: *const Complex64 = value.value as *const Complex64;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::CF64(cloned)
            }
            Tag::HList => {
                let ptr: *const Vec<Datum> = value.value as *const Vec<Datum>;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::HList(cloned)
            }
            Tag::F64Tensor => {
                let ptr: *const Tensor<f64> = value.value as *const Tensor<f64>;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::F64Tensor(cloned)
            }
            Tag::CF64Tensor => {
                let ptr: *const Tensor<Complex64> = value.value as *const Tensor<Complex64>;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::CF64Tensor(cloned)
            }
            Tag::I64Tensor => {
                let ptr: *const Tensor<i64> = value.value as *const Tensor<i64>;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::I64Tensor(cloned)
            }
            Tag::I32Tensor => {
                let ptr: *const Tensor<i32> = value.value as *const Tensor<i32>;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::I32Tensor(cloned)
            }
            Tag::Opaque => {
                let ptr: *const OpaqueData = value.value as *const OpaqueData;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::Opaque(cloned)
            }
            Tag::Function => {
                let ptr: *const Function = value.value as *const Function;
                let ptr = unsafe { Arc::from_raw(ptr) };
                let cloned = ptr.clone();
                std::mem::forget(ptr);
                Datum::Function(cloned)
            }
        }
    }

    // XXX: we don't use this to avoid Rust UB (mutable aliasing)
    // If we gain the reference from a CDatum (reference or owned doesn't matter),
    // we possible read something that is being written by another thread. The
    // writer may be gained by `Arc::get_mut`.
    unsafe fn get_moved_or_cloned_array(value: CDatum) -> Option<MoveOrClonedArray> {
        let mut cloned = false;
        let variant = match value.tag {
            Tag::I64
            | Tag::F64
            | Tag::Bool
            | Tag::U8Str
            | Tag::HList
            | Tag::CF64
            | Tag::Opaque
            | Tag::Function => return None,
            Tag::F64Tensor => {
                let ptr: *const Tensor<f64> = unsafe { std::mem::transmute(value.value) };
                let mut arc = Arc::from_raw(ptr);
                if Arc::get_mut(&mut arc).is_some() {
                    // moved
                    MoveOrClonedArrayVariant::F64Tensor(arc)
                } else {
                    cloned = true;
                    let cloned_arc = arc.clone();
                    std::mem::forget(arc);
                    MoveOrClonedArrayVariant::F64Tensor(cloned_arc)
                }
            }
            Tag::CF64Tensor => {
                let ptr: *const Tensor<Complex64> = unsafe { std::mem::transmute(value.value) };
                let mut arc = Arc::from_raw(ptr);
                if Arc::get_mut(&mut arc).is_some() {
                    // moved
                    MoveOrClonedArrayVariant::CF64Tensor(arc)
                } else {
                    cloned = true;
                    let cloned_arc = arc.clone();
                    std::mem::forget(arc);
                    MoveOrClonedArrayVariant::CF64Tensor(cloned_arc)
                }
            }
            Tag::I64Tensor => {
                let ptr: *const Tensor<i64> = unsafe { std::mem::transmute(value.value) };
                let mut arc = Arc::from_raw(ptr);
                if Arc::get_mut(&mut arc).is_some() {
                    // moved
                    MoveOrClonedArrayVariant::I64Tensor(arc)
                } else {
                    cloned = true;
                    let cloned_arc = arc.clone();
                    std::mem::forget(arc);
                    MoveOrClonedArrayVariant::I64Tensor(cloned_arc)
                }
            }
            Tag::I32Tensor => {
                let ptr: *const Tensor<i32> = unsafe { std::mem::transmute(value.value) };
                let mut arc = Arc::from_raw(ptr);
                if Arc::get_mut(&mut arc).is_some() {
                    // moved
                    MoveOrClonedArrayVariant::I32Tensor(arc)
                } else {
                    cloned = true;
                    let cloned_arc = arc.clone();
                    std::mem::forget(arc);
                    MoveOrClonedArrayVariant::I32Tensor(cloned_arc)
                }
            }
        };
        return Some(MoveOrClonedArray { cloned, variant });
    }

    #[inline]
    pub fn leak_to_reference(self) -> CDatum {
        match self {
            Datum::I64(value) => CDatum {
                value: unsafe { std::mem::transmute(value) },
                extra: 0,
                tag: Tag::I64,
            },
            Datum::F64(value) => CDatum {
                value: unsafe { std::mem::transmute(value) },
                extra: 0,
                tag: Tag::F64,
            },
            Datum::Bool(value) => CDatum {
                value: value as u64,
                extra: 0,
                tag: Tag::Bool,
            },
            Datum::CF64(value) => {
                let ptr = Arc::into_raw(value);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::CF64,
                }
            }
            Datum::U8Str(value) => {
                let ptr = Arc::into_raw(value);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::U8Str,
                }
            }
            Datum::HList(vec) => {
                let ptr = Arc::into_raw(vec);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::HList,
                }
            }
            Datum::F64Tensor(tensor) => {
                let ptr = Arc::into_raw(tensor);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::F64Tensor,
                }
            }
            Datum::CF64Tensor(tensor) => {
                let ptr = Arc::into_raw(tensor);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::CF64Tensor,
                }
            }
            Datum::I64Tensor(tensor) => {
                let ptr = Arc::into_raw(tensor);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::I64Tensor,
                }
            }
            Datum::I32Tensor(tensor) => {
                let ptr = Arc::into_raw(tensor);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::I32Tensor,
                }
            }
            Datum::Opaque(opaque) => {
                let vtable = opaque.vtable;
                let ptr = Arc::into_raw(opaque);
                CDatum {
                    value: ptr as u64,
                    extra: vtable,
                    tag: Tag::Opaque,
                }
            }
            Datum::Function(func) => {
                let ptr = Arc::into_raw(func);
                CDatum {
                    value: ptr as u64,
                    extra: 0,
                    tag: Tag::Function,
                }
            }
        }
    }
}

pub(crate) mod exports {
    use std::ops::Deref;

    use super::*;

    // ====================== MISC APIs ======================
    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_free(value: CDatum) {
        unsafe { Datum::drop_reference(value) };
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_new_runtime(out: *mut RuntimeIdent, rt: *mut CRuntime) {
        let rt_ref = &*rt;
        let name_len = c_strlen(rt_ref.name);
        let name_ptr = rt_ref.name;
        let thread_safe = rt_ref.thread_safe != 0;
        let _ = rt_ref;

        let slice = unsafe { std::slice::from_raw_parts(name_ptr, name_len) };
        let s = String::from_utf8_lossy(slice);
        let runtime_id = RUNTIMES.push(Runtime {
            name: s.into_owned(),
            thread_safe,
            def_thread: get_current_thread_id(),
        });

        *out = RuntimeIdent(runtime_id);
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_new_vtable(
        destructor: unsafe extern "C" fn(closure: *mut u8),
        runtime: RuntimeIdent,
        name: *const u8,
    ) -> u32 {
        let name_len = c_strlen(name);
        let name_slice = unsafe { std::slice::from_raw_parts(name, name_len) };
        let name = String::from_utf8_lossy(name_slice);

        let vtable = OpaqueVTable {
            name: name.into_owned(),
            foreign_destructor: destructor,
            runtime,
        };

        let vtable_id = OPAQUE_VTABLES.push(vtable);
        if vtable_id >= u32::MAX as usize {
            panic!("fatal: register too many opaque vtables");
        }

        vtable_id as u32
    }

    /// this is necessary in case we want to concurrently use tensors
    /// from runtime that is not thread-safe.
    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_run_runtime_finalizers(
        runtime: RuntimeIdent,
    ) -> ErrorMsg {
        run_finalizers(runtime)
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_free_exception(error: ErrorMsg) {
        if error.is_ok() {
            return;
        }

        error.free();
    }

    // useful when we want to create an error message without a runtime
    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_new_exception(msg: *const u8) -> ErrorMsg {
        let msg_len = c_strlen(msg);
        let msg_slice = unsafe { std::slice::from_raw_parts(msg, msg_len) };
        let msg = String::from_utf8_lossy(msg_slice);
        ErrorMsg::create_error(Cow::Owned(msg.into_owned()))
    }

    // ====================== MK APIs ======================
    // MK APIs are responsible for creating new objects.

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_opaque(
        out: *mut CDatum,
        vtable: u32,
        data: *mut u8,
    ) -> ErrorMsg {
        match OPAQUE_VTABLES.get(vtable as usize) {
            Some(_) => {}
            None => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_UNKNOWN_OPAQUE_VTABLE)),
        };

        let ptr = Arc::into_raw(Arc::new(OpaqueData {
            vtable: vtable,
            data: data,
        }));

        *out = CDatum {
            value: ptr as u64,
            extra: 0,
            tag: Tag::Opaque,
        };

        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_function(
        out: *mut CDatum,
        runtime: RuntimeIdent,
        closure: *mut u8,
        drop: unsafe extern "C" fn(closure: *mut u8),
        fptr: unsafe extern "C" fn(
            out: *mut CDatum,
            closure: *mut u8,
            args: *const CDatum,
            nargs: usize,
        ) -> ErrorMsg,
    ) -> ErrorMsg {
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_i64(out: *mut CDatum, value: i64) -> ErrorMsg {
        *out = Datum::I64(value).leak_to_reference();
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_f64(out: *mut CDatum, value: f64) -> ErrorMsg {
        *out = Datum::F64(value).leak_to_reference();
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_bool(out: *mut CDatum, value: u8) -> ErrorMsg {
        *out = Datum::Bool(if value != 0 { 1 } else { 0 }).leak_to_reference();
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_u8str(
        out: *mut CDatum,
        value: *const u8,
        len: usize,
    ) -> ErrorMsg {
        let slice = unsafe { std::slice::from_raw_parts(value, len) };
        let s = String::from_utf8_lossy(slice);
        *out = Datum::U8Str(Arc::new(s.into_owned())).leak_to_reference();
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_hlist(
        out: *mut CDatum,
        values: *const CDatum,
        len: usize,
    ) -> ErrorMsg {
        let mut vec = Vec::new();
        match vec.try_reserve(len) {
            Err(_) => {
                return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_LIST_ALLOC_OUT_OF_MEMORY))
            }
            Ok(_) => (),
        }

        for i in 0..len {
            let reference = *values.add(i);
            let datum = Datum::create_owner(reference);
            vec.push(datum);
        }

        let ptr = Arc::into_raw(Arc::new(vec));

        *out = CDatum {
            value: ptr as u64,
            extra: 0,
            tag: Tag::HList,
        };

        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_mk_tensor(
        out: *mut CDatum,
        typecode: *const u8,
        ndims: usize,
        shape: *const usize,
        strides: *const isize,
        data: *const u8,
        is_fortran_order: u8,
        drop: unsafe extern "C" fn(closure: *mut u8),
        closure: *mut u8,
        thread_safe: u8,
        runtime: RuntimeIdent,
    ) -> ErrorMsg {
        if shape.is_null() {
            return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_EMPTY_SHAPE));
        }

        let shape = unsafe { std::slice::from_raw_parts(shape, ndims) };
        let strides = if strides.is_null() {
            None
        } else {
            Some(unsafe { std::slice::from_raw_parts(strides, ndims) })
        };

        let typecode_slice: &[u8] = {
            let strlen = c_strlen(typecode);
            if strlen == 0 {
                &[]
            } else {
                unsafe { std::slice::from_raw_parts(typecode, strlen) }
            }
        };

        match typecode_slice {
            b"f64" => {
                let tensor = Tensor::adapt_native(
                    data as *const f64,
                    shape,
                    strides,
                    drop,
                    closure,
                    is_fortran_order != 0,
                    thread_safe != 0,
                    runtime,
                );
                *out = Datum::F64Tensor(Arc::new(tensor)).leak_to_reference();
            }
            b"cf64" => {
                let tensor = Tensor::adapt_native(
                    data as *const Complex64,
                    shape,
                    strides,
                    drop,
                    closure,
                    is_fortran_order != 0,
                    thread_safe != 0,
                    runtime,
                );
                *out = Datum::CF64Tensor(Arc::new(tensor)).leak_to_reference();
            }
            b"i64" => {
                let tensor = Tensor::adapt_native(
                    data as *const i64,
                    shape,
                    strides,
                    drop,
                    closure,
                    is_fortran_order != 0,
                    thread_safe != 0,
                    runtime,
                );
                *out = Datum::I64Tensor(Arc::new(tensor)).leak_to_reference();
            }
            b"i32" => {
                let tensor = Tensor::adapt_native(
                    data as *const i32,
                    shape,
                    strides,
                    drop,
                    closure,
                    is_fortran_order != 0,
                    thread_safe != 0,
                    runtime,
                );
                *out = Datum::I32Tensor(Arc::new(tensor)).leak_to_reference();
            }
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_INVALID_TYPECODE)),
        };

        ErrorMsg::ok()
    }

    // ====================== PREALLOC APIs ======================
    // Prealloc APIs are responsible for preallocating tensors.
    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_prealloc_tensor(
        out: *mut CDatum,
        typecode: *const u8,
        ndims: usize,
        shape: *const usize,
        is_fortran_order: u8,
    ) -> ErrorMsg {
        let typecode_slice: &[u8] = {
            let strlen = c_strlen(typecode);
            if strlen == 0 {
                &[]
            } else {
                unsafe { std::slice::from_raw_parts(typecode, strlen) }
            }
        };

        if shape.is_null() {
            return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_EMPTY_SHAPE));
        }
        let shape = unsafe { std::slice::from_raw_parts(shape, ndims) };

        match typecode_slice {
            b"f64" => match Tensor::prealloc(shape, is_fortran_order != 0) {
                Ok(tensor) => {
                    *out = Datum::F64Tensor(Arc::new(tensor)).leak_to_reference();
                }
                Err(_) => {
                    return ErrorMsg::create_error(Cow::Borrowed(
                        MESSAGE_TENSOR_ALLOC_OUT_OF_MEMORY,
                    ))
                }
            },
            b"cf64" => match Tensor::prealloc(shape, is_fortran_order != 0) {
                Ok(tensor) => {
                    *out = Datum::CF64Tensor(Arc::new(tensor)).leak_to_reference();
                }
                Err(_) => {
                    return ErrorMsg::create_error(Cow::Borrowed(
                        MESSAGE_TENSOR_ALLOC_OUT_OF_MEMORY,
                    ))
                }
            },
            b"i64" => match Tensor::prealloc(shape, is_fortran_order != 0) {
                Ok(tensor) => {
                    *out = Datum::I64Tensor(Arc::new(tensor)).leak_to_reference();
                }
                Err(_) => {
                    return ErrorMsg::create_error(Cow::Borrowed(
                        MESSAGE_TENSOR_ALLOC_OUT_OF_MEMORY,
                    ))
                }
            },
            b"i32" => match Tensor::prealloc(shape, is_fortran_order != 0) {
                Ok(tensor) => {
                    *out = Datum::I32Tensor(Arc::new(tensor)).leak_to_reference();
                }
                Err(_) => {
                    return ErrorMsg::create_error(Cow::Borrowed(
                        MESSAGE_TENSOR_ALLOC_OUT_OF_MEMORY,
                    ))
                }
            },
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_INVALID_TYPECODE)),
        }

        ErrorMsg::ok()
    }

    // ====================== TENSOR APIs ======================
    // TENSOR APIs are responsible for accessing and manipulating tensors.

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_readonly_element(
        out: *mut *const u8,
        reference: CDatum,
        index: usize,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        let ptr = match &datum {
            Datum::F64Tensor(tensor) => tensor
                .get_element(index)
                .map(|it| it as *const _ as *const u8),
            Datum::CF64Tensor(tensor) => tensor
                .get_element(index)
                .map(|it| it as *const _ as *const u8),
            Datum::I64Tensor(tensor) => tensor
                .get_element(index)
                .map(|it| it as *const _ as *const u8),
            Datum::I32Tensor(tensor) => tensor
                .get_element(index)
                .map(|it| it as *const _ as *const u8),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)),
        };

        match ptr {
            Some(ptr) => unsafe { *out = ptr },
            None => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_OUT_OF_BOUNDS)),
        }

        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_update_element(
        in_and_out_ref: *mut CDatum,
        index: usize,
        value_ref: *const u8,
    ) -> ErrorMsg {
        if in_and_out_ref.is_null() || value_ref.is_null() {
            return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NULL_REFERENCE));
        }
        let original = std::ptr::read(in_and_out_ref);
        let MoveOrClonedArray {
            cloned,
            mut variant,
        } = match Datum::get_moved_or_cloned_array(original)
            .ok_or_else(|| ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)))
        {
            Ok(it) => it,
            Err(err) => return err,
        };

        match &mut variant {
            MoveOrClonedArrayVariant::F64Tensor(tensor) => {
                if index >= tensor.numel() {
                    if !cloned {
                        // if moved and raise,
                        // we should not drop the original array
                        std::mem::forget(variant);
                    }
                    return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_OUT_OF_BOUNDS));
                }
                let it = Arc::make_mut(tensor);
                let lval = it.get_element_mut(index).expect("index is already checked");
                *lval = std::ptr::read(value_ref as *const f64);
            }
            MoveOrClonedArrayVariant::CF64Tensor(tensor) => {
                if index >= tensor.numel() {
                    if !cloned {
                        // if moved and raise,
                        // we should not drop the original array
                        std::mem::forget(variant);
                    }
                    return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_OUT_OF_BOUNDS));
                }
                let it = Arc::make_mut(tensor);
                let lval = it.get_element_mut(index).expect("index is already checked");
                *lval = std::ptr::read(value_ref as *const Complex64);
            }
            MoveOrClonedArrayVariant::I64Tensor(tensor) => {
                if index >= tensor.numel() {
                    if !cloned {
                        // if moved and raise,
                        // we should not drop the original array
                        std::mem::forget(variant);
                    }
                    return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_OUT_OF_BOUNDS));
                }
                let it = Arc::make_mut(tensor);
                let lval = it.get_element_mut(index).expect("index is already checked");
                *lval = std::ptr::read(value_ref as *const i64);
            }
            MoveOrClonedArrayVariant::I32Tensor(tensor) => {
                if index >= tensor.numel() {
                    if !cloned {
                        // if moved and raise,
                        // we should not drop the original array
                        std::mem::forget(variant);
                    }
                    return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_OUT_OF_BOUNDS));
                }
                let it = Arc::make_mut(tensor);
                let lval = it.get_element_mut(index).expect("index is already checked");
                *lval = std::ptr::read(value_ref as *const i32)
            }
        };

        if !cloned {
            // if moved, we should not drop the original array
            // but still return it (do not change reference)
        } else {
            // if cloned, the original is different from 'variant'.
            // when the update succeeded, the original array
            // has been replaced with the updated one, so that
            // we drop the previous reference.

            *in_and_out_ref = match variant {
                MoveOrClonedArrayVariant::F64Tensor(tensor) => {
                    Datum::F64Tensor(tensor).leak_to_reference()
                }
                MoveOrClonedArrayVariant::CF64Tensor(tensor) => {
                    Datum::CF64Tensor(tensor).leak_to_reference()
                }
                MoveOrClonedArrayVariant::I64Tensor(tensor) => {
                    Datum::I64Tensor(tensor).leak_to_reference()
                }
                MoveOrClonedArrayVariant::I32Tensor(tensor) => {
                    Datum::I32Tensor(tensor).leak_to_reference()
                }
            };

            // drop the original array
            tonglun_free(original);
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_readonly_data(
        out: *mut *const u8,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        let ptr = match &datum {
            Datum::F64Tensor(tensor) => tensor.get_data_ptr() as *const u8,
            Datum::CF64Tensor(tensor) => tensor.get_data_ptr() as *const u8,
            Datum::I64Tensor(tensor) => tensor.get_data_ptr() as *const u8,
            Datum::I32Tensor(tensor) => tensor.get_data_ptr() as *const u8,
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)),
        };

        unsafe { *out = ptr };
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_shape(
        out: *mut *const usize,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        let shape = match &datum {
            Datum::F64Tensor(tensor) => tensor.shape(),
            Datum::CF64Tensor(tensor) => tensor.shape(),
            Datum::I64Tensor(tensor) => tensor.shape(),
            Datum::I32Tensor(tensor) => tensor.shape(),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)),
        };

        unsafe { *out = shape.as_ptr() };
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_strides(
        out: *mut *const isize,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        let strides = match &datum {
            Datum::F64Tensor(tensor) => tensor.strides(),
            Datum::CF64Tensor(tensor) => tensor.strides(),
            Datum::I64Tensor(tensor) => tensor.strides(),
            Datum::I32Tensor(tensor) => tensor.strides(),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)),
        };
        unsafe { *out = strides.as_ptr() };
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_numel(
        out: *mut usize,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::F64Tensor(tensor) => *out = tensor.numel(),
            Datum::CF64Tensor(tensor) => *out = tensor.numel(),
            Datum::I64Tensor(tensor) => *out = tensor.numel(),
            Datum::I32Tensor(tensor) => *out = tensor.numel(),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)),
        }

        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_tensor_size(
        out: *mut usize,
        reference: CDatum,
        axis: usize,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::F64Tensor(tensor) => *out = tensor.size(axis),
            Datum::CF64Tensor(tensor) => *out = tensor.size(axis),
            Datum::I64Tensor(tensor) => *out = tensor.size(axis),
            Datum::I32Tensor(tensor) => *out = tensor.size(axis),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_TENSOR)),
        }

        ErrorMsg::ok()
    }

    // ====================== Ops ======================
    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_i64_load(
        out: *mut i64,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::I64(value) => *out = *value,
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_I64)),
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_f64_load(
        out: *mut f64,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::F64(value) => *out = *value,
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_F64)),
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_complex64_load(
        out: *mut Complex64,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::CF64(value) => *out = *value.deref(),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_CF64)),
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_bool_load(
        out: *mut u8,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::Bool(value) => *out = *value,
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_BOOL)),
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_str_load(
        out: *mut *const u8,
        len: *mut usize,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::U8Str(value) => {
                *out = value.as_ptr();
                if !len.is_null() {
                    *len = value.len();
                }
            }
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_U8STR)),
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_list_loadi(
        out: *mut CDatum,
        reference: CDatum,
        index: usize,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::HList(elements) => {
                if index >= elements.len() {
                    return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_OUT_OF_BOUNDS));
                }
                let elt = elements[index].clone().leak_to_reference();
                *out = elt;
            }
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_HLIST)),
        }
        ErrorMsg::ok()
    }

    #[no_mangle]
    pub(crate) unsafe extern "C" fn tonglun_ops_list_len(
        out: *mut usize,
        reference: CDatum,
    ) -> ErrorMsg {
        let datum = Datum::create_owner(reference);
        match &datum {
            Datum::HList(elements) => *out = elements.len(),
            _ => return ErrorMsg::create_error(Cow::Borrowed(MESSAGE_NOT_HLIST)),
        }

        ErrorMsg::ok()
    }
}

fn c_strlen(s: *const u8) -> usize {
    let mut len = 0;
    while unsafe { *s.add(len) } != 0 {
        len += 1;
    }
    len
}

#[allow(unused)]
#[cfg(test)]
mod tests {
    use super::*;
    use std::{ptr, sync::atomic::AtomicBool};

    // Helper function to create a test runtime
    fn create_test_runtime() -> RuntimeIdent {
        let mut runtime = RuntimeIdent(0);
        let mut crt = CRuntime {
            name: b"test_runtime\0".as_ptr(),
            thread_safe: 1,
        };
        unsafe {
            exports::tonglun_new_runtime(
                &mut runtime as *mut RuntimeIdent,
                &mut crt as *mut CRuntime,
            );
        }
        runtime
    }

    // Test tensor creation and manipulation
    #[test]
    fn test_tensor_operations() {
        // Test F64 tensor creation
        let shape = vec![2, 3];
        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };
        let result = unsafe {
            exports::tonglun_prealloc_tensor(
                &mut tensor_datum as *mut CDatum,
                b"f64\0".as_ptr(),
                2,
                shape.as_ptr(),
                0,
            )
        };
        assert!(result.is_ok());

        // Test tensor shape access
        let mut shape_ptr: *const usize = ptr::null();
        let result = unsafe {
            exports::tonglun_tensor_shape(&mut shape_ptr as *mut *const usize, tensor_datum)
        };
        assert!(result.is_ok());

        let retrieved_shape = unsafe { std::slice::from_raw_parts(shape_ptr, 2) };
        assert_eq!(retrieved_shape, &[2, 3]);

        // Cleanup
        unsafe { exports::tonglun_free(tensor_datum) };
    }

    // Test Datum creation and conversion
    #[test]
    fn test_datum_creation() {
        // Test i64 creation
        let mut i64_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::I64,
        };
        let result = unsafe { exports::tonglun_mk_i64(&mut i64_datum as *mut CDatum, 42) };
        assert!(result.is_ok());

        // Test i64 loading
        let mut value: i64 = 0;
        let result = unsafe { exports::tonglun_ops_i64_load(&mut value as *mut i64, i64_datum) };
        assert!(result.is_ok());
        assert_eq!(value, 42);

        // Cleanup
        unsafe { exports::tonglun_free(i64_datum) };
    }

    // Test string handling
    #[test]
    fn test_string_operations() {
        let test_str = b"hello world";
        let mut str_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::U8Str,
        };
        let result = unsafe {
            exports::tonglun_mk_u8str(
                &mut str_datum as *mut CDatum,
                test_str.as_ptr(),
                test_str.len(),
            )
        };
        assert!(result.is_ok());

        // Test string loading
        let mut str_ptr: *const u8 = ptr::null();
        let mut len: usize = 0;
        let result = unsafe {
            exports::tonglun_ops_str_load(
                &mut str_ptr as *mut *const u8,
                &mut len as *mut usize,
                str_datum,
            )
        };
        assert!(result.is_ok());
        assert_eq!(len, test_str.len());

        let loaded_str = unsafe { std::slice::from_raw_parts(str_ptr, len) };
        assert_eq!(loaded_str, test_str);

        // Cleanup
        unsafe { exports::tonglun_free(str_datum) };
    }

    // Test list operations
    #[test]
    fn test_list_operations() {
        // Create test elements
        let mut i64_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::I64,
        };
        unsafe { exports::tonglun_mk_i64(&mut i64_datum, 42) };

        let elements = vec![i64_datum];
        let mut list_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::HList,
        };
        let result = unsafe {
            exports::tonglun_mk_hlist(
                &mut list_datum as *mut CDatum,
                elements.as_ptr(),
                elements.len(),
            )
        };
        assert!(result.is_ok());

        // Test list length
        let mut len: usize = 0;
        let result = unsafe { exports::tonglun_ops_list_len(&mut len as *mut usize, list_datum) };
        assert!(result.is_ok());
        assert_eq!(len, 1);

        // Test list element access
        let mut element_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::I64,
        };
        let result = unsafe {
            exports::tonglun_ops_list_loadi(&mut element_datum as *mut CDatum, list_datum, 0)
        };
        assert!(result.is_ok());

        let mut value: i64 = 0;
        unsafe {
            exports::tonglun_ops_i64_load(&mut value as *mut i64, element_datum);
        }
        assert_eq!(value, 42);

        // Cleanup
        unsafe {
            exports::tonglun_free(element_datum);
            exports::tonglun_free(list_datum);
        }
    }

    // Test tensor data access
    #[test]
    fn test_tensor_data_access() {
        let runtime = create_test_runtime();
        let shape = vec![2, 2];
        let data = vec![1.0f64, 2.0, 3.0, 4.0];

        unsafe extern "C" fn test_drop(_: *mut u8) {}

        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };
        let result = unsafe {
            exports::tonglun_mk_tensor(
                &mut tensor_datum,
                b"f64\0".as_ptr(),
                2,
                shape.as_ptr(),
                ptr::null(),
                data.as_ptr() as *const u8,
                0, // is_fortran_order
                test_drop,
                ptr::null_mut(),
                1, // thread_safe
                runtime,
            )
        };
        assert!(result.is_ok());

        // Test readonly data access
        let mut data_ptr: *const u8 = ptr::null();
        let result = unsafe {
            exports::tonglun_tensor_readonly_data(&mut data_ptr as *mut *const u8, tensor_datum)
        };
        assert!(result.is_ok());
        assert!(!data_ptr.is_null());

        // Cleanup
        unsafe { exports::tonglun_free(tensor_datum) };
    }

    // Test complex number operations
    #[test]
    fn test_complex_operations() {
        let mut cf64_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::CF64,
        };

        // Create a complex number through Datum first
        let complex = Complex64::new(1.0, 2.0);
        let datum = Datum::CF64(Arc::new(complex));
        cf64_datum = datum.leak_to_reference();

        // Test complex loading
        let mut value = Complex64::new(0.0, 0.0);
        let result = unsafe {
            exports::tonglun_ops_complex64_load(&mut value as *mut Complex64, cf64_datum)
        };
        assert!(result.is_ok());
        assert_eq!(value, complex);

        // Cleanup
        unsafe { exports::tonglun_free(cf64_datum) };
    }

    // Test boolean operations
    #[test]
    fn test_boolean_operations() {
        let mut bool_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::Bool,
        };

        let result = unsafe { exports::tonglun_mk_bool(&mut bool_datum as *mut CDatum, 1) };
        assert!(result.is_ok());

        // Test boolean loading
        let mut value: u8 = 0;
        let result = unsafe { exports::tonglun_ops_bool_load(&mut value as *mut u8, bool_datum) };
        assert!(result.is_ok());
        assert_eq!(value, 1);

        // Cleanup
        unsafe { exports::tonglun_free(bool_datum) };
    }

    // Test f64 operations
    #[test]
    fn test_f64_operations() {
        let mut f64_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64,
        };

        let result = unsafe { exports::tonglun_mk_f64(&mut f64_datum as *mut CDatum, 3.14159) };
        assert!(result.is_ok());

        // Test f64 loading
        let mut value: f64 = 0.0;
        let result = unsafe { exports::tonglun_ops_f64_load(&mut value as *mut f64, f64_datum) };
        assert!(result.is_ok());
        assert!((value - 3.14159).abs() < f64::EPSILON);

        // Cleanup
        unsafe { exports::tonglun_free(f64_datum) };
    }

    // Test tensor element access
    #[test]
    fn test_tensor_element_access() {
        let runtime = create_test_runtime();
        let shape = vec![2, 2];
        let data = vec![1.0f64, 2.0, 3.0, 4.0];

        unsafe extern "C" fn test_drop(_: *mut u8) {}

        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };

        let result = unsafe {
            exports::tonglun_mk_tensor(
                &mut tensor_datum,
                b"f64\0".as_ptr(),
                2,
                shape.as_ptr(),
                ptr::null(),
                data.as_ptr() as *const u8,
                0,
                test_drop,
                ptr::null_mut(),
                1,
                runtime,
            )
        };
        assert!(result.is_ok());

        // Test readonly element access
        let mut element_ptr: *const u8 = ptr::null();
        let result = unsafe {
            exports::tonglun_tensor_readonly_element(
                &mut element_ptr as *mut *const u8,
                tensor_datum,
                0,
            )
        };
        assert!(result.is_ok());
        assert!(!element_ptr.is_null());

        let mut numel: usize = 0;
        let result =
            unsafe { exports::tonglun_tensor_numel(&mut numel as *mut usize, tensor_datum) };
        assert!(result.is_ok());
        assert_eq!(numel, 4);

        let elements = unsafe { std::slice::from_raw_parts(element_ptr, numel) };
        println!("elements: {:?}", elements);

        let value = unsafe { *(element_ptr as *const f64) };
        assert_eq!(value, 1.0);

        // Test tensor size
        let mut size: usize = 0;
        let result =
            unsafe { exports::tonglun_tensor_size(&mut size as *mut usize, tensor_datum, 0) };
        assert!(result.is_ok());
        assert_eq!(size, 2);

        // Test tensor numel
        let mut numel: usize = 0;
        let result =
            unsafe { exports::tonglun_tensor_numel(&mut numel as *mut usize, tensor_datum) };
        assert!(result.is_ok());
        assert_eq!(numel, 4);

        // Test tensor strides
        let mut strides_ptr: *const isize = ptr::null();
        let result = unsafe {
            exports::tonglun_tensor_strides(&mut strides_ptr as *mut *const isize, tensor_datum)
        };
        assert!(result.is_ok());
        assert!(!strides_ptr.is_null());

        // Cleanup
        unsafe { exports::tonglun_free(tensor_datum) };
    }

    // Test error handling
    #[test]
    fn test_error_handling() {
        // Test invalid tensor type code
        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };
        let shape = vec![2, 2];
        let result = unsafe {
            unsafe extern "C" fn drop_fn(it: *mut u8) {
                println!("Dropping tensor with closure {:?}", it);
            }
            exports::tonglun_mk_tensor(
                &mut tensor_datum,
                b"invalid\0".as_ptr(),
                2,
                shape.as_ptr(),
                ptr::null(),
                ptr::null(),
                0,
                drop_fn,
                ptr::null_mut(),
                1,
                create_test_runtime(),
            )
        };

        // Test out of bounds access
        let mut i64_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::I64,
        };
        unsafe { exports::tonglun_mk_i64(&mut i64_datum, 42) };
        let elements = vec![i64_datum];
        let mut list_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::HList,
        };
        unsafe {
            exports::tonglun_mk_hlist(
                &mut list_datum as *mut CDatum,
                elements.as_ptr(),
                elements.len(),
            )
        };

        let mut element_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::I64,
        };
        let result = unsafe {
            exports::tonglun_ops_list_loadi(&mut element_datum as *mut CDatum, list_datum, 999)
        };
        assert_eq!(
            result.clone_error_msg(),
            Some(MESSAGE_OUT_OF_BOUNDS.to_string())
        );

        // Cleanup
        unsafe {
            exports::tonglun_free(i64_datum);
            exports::tonglun_free(list_datum);
        };
    }

    // Test runtime finalizers
    #[test]
    fn test_runtime_finalizers() {
        let runtime = create_test_runtime();
        println!("\nruntime: {:?}\n", runtime);

        // Create a tensor with a destructor
        let mut drop_called = AtomicBool::new(false);

        unsafe extern "C" fn test_drop(closure: *mut u8) {
            println!("test_drop: {:?}", closure);
            let drop_called = closure as *mut AtomicBool;
            let drop_called = drop_called.as_ref().unwrap();
            drop_called.store(true, std::sync::atomic::Ordering::SeqCst);
            // let mut it = (&*drop_called).borrow_mut();
            // *it = true;
            // drop(it);
        }

        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };
        let shape = vec![1];
        let data = vec![1.0f64];

        unsafe {
            exports::tonglun_mk_tensor(
                &mut tensor_datum,
                b"f64\0".as_ptr(),
                1,
                shape.as_ptr(),
                ptr::null(),
                data.as_ptr() as *const u8,
                0,
                test_drop,
                (&drop_called) as *const _ as *mut u8,
                0, // not thread safe
                runtime,
            )
        };

        // Run finalizers and verify the destructor was called
        unsafe { exports::tonglun_free(tensor_datum) };
        unsafe { exports::tonglun_run_runtime_finalizers(runtime) };
        assert!(drop_called.load(std::sync::atomic::Ordering::SeqCst));
    }

    #[test]
    fn test_update_element() {
        let runtime = create_test_runtime();
        let shape = vec![2, 2];
        let data = vec![1.0f64, 2.0, 3.0, 4.0];

        unsafe extern "C" fn test_drop(_: *mut u8) {
            println!("[test_update_element] Dropping tensor");
        }

        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };

        let result = unsafe {
            exports::tonglun_mk_tensor(
                &mut tensor_datum,
                b"f64\0".as_ptr(),
                2,
                shape.as_ptr(),
                ptr::null(), // strides
                data.as_ptr() as *const u8,
                0,
                test_drop,
                ptr::null_mut(),
                1,
                runtime,
            )
        };
        assert!(result.is_ok());

        // Test readonly element access
        let mut element_ptr: *const u8 = ptr::null();
        let result = unsafe {
            exports::tonglun_tensor_readonly_element(
                &mut element_ptr as *mut *const u8,
                tensor_datum,
                0,
            )
        };
        assert!(result.is_ok());
        assert!(!element_ptr.is_null());

        let mut numel: usize = 0;
        let result =
            unsafe { exports::tonglun_tensor_numel(&mut numel as *mut usize, tensor_datum) };
        assert!(result.is_ok());
        assert_eq!(numel, 4);

        let elements = unsafe { std::slice::from_raw_parts(element_ptr, numel) };
        println!("elements: {:?}", elements);

        let value = unsafe { *(element_ptr as *const f64) };
        assert_eq!(value, 1.0);

        // Test update element
        let new_value = 42.0;
        let prev = tensor_datum;
        let result = unsafe {
            exports::tonglun_tensor_update_element(
                &mut tensor_datum,
                0,
                &new_value as *const f64 as *const u8,
            )
        };
        assert!(result.is_ok());
        let after = tensor_datum;

        // ensure no cow happened
        assert_eq!(prev, after);

        // Test readonly element access
        let mut element_ptr: *const u8 = ptr::null();
        let result = unsafe {
            exports::tonglun_tensor_readonly_element(
                &mut element_ptr as *mut *const u8,
                tensor_datum,
                0,
            )
        };

        let value = unsafe { *(element_ptr as *const f64) };
        assert_eq!(value, 42.0);
    }

    #[test]
    fn test_update_element_cow() {
        let runtime = create_test_runtime();
        let shape = vec![2, 2];
        let data = vec![1.0f64, 2.0, 3.0, 4.0];

        unsafe extern "C" fn test_drop(_: *mut u8) {
            println!("[test_update_element_cow] Dropping tensor");
        }

        let mut tensor_datum = CDatum {
            value: 0,
            extra: 0,
            tag: Tag::F64Tensor,
        };

        let result = unsafe {
            exports::tonglun_mk_tensor(
                &mut tensor_datum,
                b"f64\0".as_ptr(),
                2,
                shape.as_ptr(),
                ptr::null(), // strides
                data.as_ptr() as *const u8,
                0,
                test_drop,
                ptr::null_mut(),
                1,
                runtime,
            )
        };
        assert!(result.is_ok());

        // if exception happened, nothing
        // nothing is changed.
        {
            // Test update element
            let new_value = 42.0;
            let prev = tensor_datum;
            let result = unsafe {
                exports::tonglun_tensor_update_element(
                    &mut tensor_datum,
                    1000,
                    &new_value as *const f64 as *const u8,
                )
            };
            assert!(!result.is_ok());
            let after = tensor_datum;

            assert_eq!(prev, after);
        }

        // lock data:
        let mut locker = Datum::create_owner(tensor_datum);

        // Test update element
        let new_value = 42.0;
        let prev = tensor_datum;
        let result = unsafe {
            exports::tonglun_tensor_update_element(
                &mut tensor_datum,
                0,
                &new_value as *const f64 as *const u8,
            )
        };

        assert!(result.is_ok());
        let after = tensor_datum;

        // ensure cow happened
        assert_ne!(prev, after);

        drop(locker);

        // Test update element
        let new_value = 42.0;
        let prev = tensor_datum;
        let result = unsafe {
            exports::tonglun_tensor_update_element(
                &mut tensor_datum,
                0,
                &new_value as *const f64 as *const u8,
            )
        };

        assert!(result.is_ok());
        let after = tensor_datum;

        // ensure cow happened
        assert_eq!(prev, after);
    }
}
